'''
Created on Jul 16, 2012

@author: nzhao
'''
import numpy as np

from traits.api import HasTraits, Instance
from math_qm.DensityMatrix import DensityMatrix
from math_qm.SpinHamiltonian import Hamiltonian

class AbstractSpinSystem(HasTraits):

    state = Instance(DensityMatrix)
    hamiltonian = Instance(Hamiltonian)

    def calculate_eigs(self):
        self.hamiltonian.calculate_eigs()
        return [self.hamiltonian.eig_vals, self.hamiltonian.eig_vect]
    
    def mr_transitions(self, mr_op, unit = 2.0*np.pi*1000.0):
        freqlist, wavefunc = self.calculate_eigs()
        idx_lst  = range(self.hamiltonian.dim)
        
        trans = [ ((i,j), 
                   (freqlist[j] - freqlist[i])/unit, 
                   np.abs( np.dot(np.conjugate(wavefunc[:, i]), np.dot(mr_op, wavefunc[:, j])))**2 
                   ) 
                 for i in idx_lst for j in idx_lst if i < j]
        return trans
    
    def mr_spectrum(self, flist, width, mr_op, shape='Gaussian', unit = 2.0*np.pi*1000.0):
        trans = self.mr_transitions(mr_op, unit)
        
        spect = []
        idx   = ['tot']
        for tr in trans:
            idx.append( str(tr[0]) )
            center = tr[1]
            amp    = tr[2]
            s = amp * self._line_shape(flist, center, width, shape)
            spect.append( s )
        spect.insert(0, np.sum(spect, axis=0))
        return [flist, spect, idx]

    def _line_shape(self, f, center, width, shape):
        res = 0.0
        if shape is 'Gaussian':
            res = np.exp(-(f-center)**2 / 2.0 / width**2 )
        elif shape is 'Lorentzian':
            res = width / (width**2 + (f-center)**2)
        else:
            print 'unknown line shape'
        
        return res